package com.dz.risk.config;

import lombok.extern.slf4j.Slf4j;
import org.apache.commons.lang3.StringUtils;
import org.springframework.aop.interceptor.AsyncUncaughtExceptionHandler;
import org.springframework.aop.interceptor.SimpleAsyncUncaughtExceptionHandler;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.context.annotation.Lazy;
import org.springframework.context.annotation.Scope;
import org.springframework.core.task.TaskDecorator;
import org.springframework.scheduling.annotation.AsyncConfigurerSupport;
import org.springframework.scheduling.concurrent.ThreadPoolTaskExecutor;
import org.springframework.web.context.request.RequestAttributes;
import org.springframework.web.context.request.RequestContextHolder;

import java.util.concurrent.RejectedExecutionHandler;
import java.util.concurrent.ThreadPoolExecutor;

/**
 * @author Doke
 */
@Configuration
@Slf4j
public class ThreadPoolConfiguration extends AsyncConfigurerSupport {

    private static final int CORE_POOL_SIZE = 20;

    private static final int QUEUE_CAPACITY = 5;

    private static final int MAX_POOL_SIZE = 1000;

    private static final int KEEP_ALIVE_SECONDS = 120;

    private static final boolean DAEMON = false;

    private static final String THREAD_NAME_PREFIX = "springWordThread";

    private static final String REJECTED_EXECUTION_HANDLER_CLASS_NAME = "java.util.concurrent.ThreadPoolExecutor$CallerRunsPolicy";

    @Bean("customizeExecutor")
    protected ThreadPoolTaskExecutor getTaskExecutor() {
        ThreadPoolTaskExecutor threadPoolTaskExecutor = new ThreadPoolTaskExecutor();
        threadPoolTaskExecutor.setCorePoolSize(CORE_POOL_SIZE);
        threadPoolTaskExecutor.setQueueCapacity(QUEUE_CAPACITY);
        threadPoolTaskExecutor.setMaxPoolSize(MAX_POOL_SIZE);
        threadPoolTaskExecutor.setKeepAliveSeconds(KEEP_ALIVE_SECONDS);
        threadPoolTaskExecutor.setDaemon(DAEMON);
        threadPoolTaskExecutor.setThreadNamePrefix(THREAD_NAME_PREFIX);
        threadPoolTaskExecutor.setTaskDecorator(new ContextDecorator());
        threadPoolTaskExecutor.setRejectedExecutionHandler(newInstanceExecutionHandler());
        return threadPoolTaskExecutor;
    }

    /**
     * 构造线程池处理拒绝策略 ，默认为CallerRunsPolicy策略
     *
     * @return RejectedExecutionHandler
     */
    private RejectedExecutionHandler newInstanceExecutionHandler() {
        //默认处理策略
        RejectedExecutionHandler rejectedExecutionHandler = new ThreadPoolExecutor.CallerRunsPolicy();
        try {
            if (!StringUtils.isEmpty(REJECTED_EXECUTION_HANDLER_CLASS_NAME)) {
                Class<?> hander = Class.forName(REJECTED_EXECUTION_HANDLER_CLASS_NAME);
                Object object = hander.newInstance();
                if (object instanceof RejectedExecutionHandler) {
                    rejectedExecutionHandler = (RejectedExecutionHandler) object;
                }
            }
        } catch (Exception e) {
            log.error(e.getMessage(), e);
        }
        return rejectedExecutionHandler;
    }


    @Override
    public AsyncUncaughtExceptionHandler getAsyncUncaughtExceptionHandler() {
        return new SimpleAsyncUncaughtExceptionHandler();
    }

    static class ContextDecorator implements TaskDecorator {
        @Override
        public Runnable decorate(Runnable runnable) {
            RequestAttributes context = RequestContextHolder.currentRequestAttributes();
            return () -> {
                try {
                    RequestContextHolder.setRequestAttributes(context);
                    runnable.run();
                } finally {
                    RequestContextHolder.resetRequestAttributes();
                }
            };
        }
    }
}
